import gzip
import io


# 转为字节：每8位一组
def bin_str_to_bytes(s:str):
    s = s.ljust((len(s) + 7) // 8 * 8, '0')  # 补齐到8的倍数
    return bytes([int(s[i:i+8], 2) for i in range(0, len(s), 8)])

# 转为二进制字符串：将字节流还原为 '0' 和 '1' 组成的字符串
def bytes_to_bin_str(data: bytes) -> str:
    return ''.join(f'{byte:08b}' for byte in data)

   # 压缩
def compressMethod(lsbs:list[str])->bytes:

    lsbsStr = ''.join(lsbs)  # 拼接成字符串
    # 转为字节
    data = bin_str_to_bytes(lsbsStr)
    # 使用 gzip 在内存中压缩
    with io.BytesIO() as buffer:
        with gzip.GzipFile(fileobj=buffer, mode='wb') as gz_file:
            gz_file.write(data)
        compressed = buffer.getvalue()

    # print(f"gzip 压缩前: {len(data)} 字节")
    # print(f"gzip 压缩后: {len(compressed)} 字节")
    # print(f"压缩比: {len(compressed)/len(data):.2f}")     
    return compressed

# 解压：将压缩数据byte解压还原为bit列表（list[str]）
def decompressMethod(compressed: bytes) -> list[str]:
    try:
        # 使用 gzip 解压
        with io.BytesIO(compressed) as buffer:
            with gzip.GzipFile(fileobj=buffer, mode='rb') as gz_file:
                decompressed_data = gz_file.read()  # 得到原始 bytes

        # 将 bytes 转回二进制字符串
        bin_str = bytes_to_bin_str(decompressed_data)

        # 转为 list['0', '1', ...]
        return list(bin_str)
    
    except Exception as e:
        raise ValueError(f"解压失败：可能数据损坏或非 gzip 格式。错误: {e}")
    
if __name__ == "__main__":
    # 示例：0 占 80%，1 占 20%
    lsbs = ["0"] * 800 + ["1"] * 200
    print(f"原始数据长度: {len(lsbs)} bits")

    # 压缩
    compressed = compressMethod(lsbs)
    print(f"压缩后数据长度: {len(compressed)} bytes")

    # 解压
    decompressed = decompressMethod(compressed)
    print(f"解压后数据长度: {len(decompressed)} bits")

    # 验证一致性
    if lsbs == decompressed:
        print("压缩 -> 解压 成功，数据一致！")
    else:
        print("数据不一致！")

    # 可选：查看压缩率
    original_bit_length = len(lsbs)
    compressed_byte_length = len(compressed)
    compression_ratio = compressed_byte_length / (original_bit_length / 8)  # 按字节比较
    print(f"压缩率: {compression_ratio:.2f} ({compressed_byte_length}/{original_bit_length//8})")